package com.jboltai.event.ai.aichat;

import cn.hutool.core.util.StrUtil;
import com.jboltai.capability.message.AIMessage;
import com.jboltai.capability.message.AIMessageRole;
import com.jboltai.capability.message.MessageFinishReason;
import com.jboltai.capability.model.AIEventProcessor;
import com.jboltai.event.*;
import com.jboltai.event.ai.AIEvent;
import com.jboltai.event.handler.AIChatEventHandler;
import com.jboltai.resource.ai.AIModel;
import com.jboltai.resource.ai.AIResource;

import java.util.Objects;

/**
 * ai聊天事件
 */
public class AIChatEvent extends AIEvent<AIChatEvent, AIMessage, AIChatEventHandler> {


    public AIChatEvent() {
        super(EventType.AI_CHAT,null, null);
    }

    public AIChatEvent(AIModel model) {
        super(EventType.AI_CHAT,model, model.getDefaultModel());
    }


    /**
     * 创建一个ai事件
     * @param model 使用哪个模型，具体的模型名称根据该枚举值的defaultModel获取
     * @param successHandler
     * @param failHandler
     */
    public AIChatEvent(AIModel model, SuccessHandler<AIChatEvent, AIMessage> successHandler, FailHandler failHandler) {
        super(EventType.AI_CHAT,model, model.getDefaultModel());
        this.successHandler = successHandler;
        this.failHandler = failHandler;
    }
    /**
     * 创建一个ai事件
     * @param model 使用哪个模型，具体的模型名称根据该枚举值的defaultModel获取
     * @param successHandler
     * @param failHandler
     */
    public AIChatEvent(AIModel model, SuccessHandler<AIChatEvent, AIMessage> successHandler, FailHandler failHandler, CompleteHandler completeHandler) {
        super(EventType.AI_CHAT,model, model.getDefaultModel());
        this.successHandler = successHandler;
        this.failHandler = failHandler;
        this.completeHandler = completeHandler;
    }



    private ThinkHandler<AIChatEvent> thinkHandler;

    private ThinkStatus thinkStatus;

    /**
     * 思考开始时间
     */
    private long thinkStartTime;
    /**
     * 思考结束时间
     */
    private long thinkEndTime;


    /**
     * 是否是流式请求
     */
    private boolean stream = true;
    /**
     * 完结原因
     */
    private MessageFinishReason finishReason;

    /**
     * 提示词token数
     */
    private Integer promptTokenCount;

    /**
     * 响应消耗的token数
     */
    private Integer completionTokenCount;

    /**
     * 总消耗的token数
     */
    private Integer totalTokenCount;

    /**
     * 最大token数
     */
    private Integer maxTokens = 2000;

    /**
     * 思考模式预算的token数
     */
    private Integer thinkingBudget = 4096;


    /**
     * AI大模型深度思考的回调
     * @param thinkHandler
     * @return
     */
    public AIChatEvent onThinking(ThinkHandler<AIChatEvent> thinkHandler) {
        this.thinkHandler = thinkHandler;
        return this;
    }

    @Override
    protected void innerPrevHandle() {
        //非流式的，思考开始时间就是事件开始的时间
        this.thinkStartTime = System.currentTimeMillis();
    }

    /**
     * 触发深度思考
     * @param think
     */
    public void thinking(String think) {
        if (this.getState() != EventState.PROCESSING ||Objects.equals(this.thinkingEnable, true) == false) {
            return;
        }
        appendThink(think);
        if (thinkStatus == null) {
            thinkStatus = ThinkStatus.START;
            //流式的会重新更新思考的开始时间，只有接收到think内容时，才算开始思考
            thinkStartTime = System.currentTimeMillis();
        } else if (thinkStatus == ThinkStatus.START) {
            thinkStatus = ThinkStatus.RUNNING;
        }
        if (thinkHandler != null) {
            thinkHandler.accept(this, think, thinkStatus);
        }
    }

    /**
     * 结束深度思考
     * @param think
     */
    public void thinkOver(String think) {
        if (this.getState() != EventState.PROCESSING ||Objects.equals(this.thinkingEnable, true) == false) {
            return;
        }
        appendThink(think);
        thinkStatus = ThinkStatus.COMPLETE;
        thinkEndTime = System.currentTimeMillis();
        if (thinkHandler != null) {
            thinkHandler.accept(this, think, ThinkStatus.COMPLETE);
        }
    }

    /**
     * 获取思考结果
     * @return
     */
    public String getThinkResult() {
        return result != null ? result.getThink() : null;
    }

    /**
     * 获取是否是stream请求，如果使用了functioncall ,会强制为false
     * @return
     */
    public boolean getIsStream() {
        return stream;
    }

    public AIChatEvent setStream(boolean stream) {
        this.stream = stream;
        return this;
    }

    @Override
    public String preCheck() {
        if (StrUtil.isBlank(this.prompt) &&
                (this.contextMessages == null || this.contextMessages.isEmpty() || this.contextMessages.stream().noneMatch(item -> item.getRole() == AIMessageRole.user))) {
            return "未设置用户消息";
        }
        return null;
    }

    public MessageFinishReason getFinishReason() {
        return finishReason;
    }

    public void setFinishReason(MessageFinishReason finishReason) {
        this.finishReason = finishReason;
    }

    public Integer getPromptTokenCount() {
        return promptTokenCount;
    }

    public void setPromptTokenCount(Integer promptTokenCount) {
        this.promptTokenCount = promptTokenCount;
    }

    public Integer getCompletionTokenCount() {
        return completionTokenCount;
    }

    public void setCompletionTokenCount(Integer completionTokenCount) {
        this.completionTokenCount = completionTokenCount;
    }

    public Integer getTotalTokenCount() {
        return totalTokenCount;
    }

    public void setTotalTokenCount(Integer totalTokenCount) {
        this.totalTokenCount = totalTokenCount;
    }

    public Integer getMaxTokens() {
        return maxTokens;
    }

    public AIChatEvent setMaxTokens(Integer maxTokens) {
        this.maxTokens = maxTokens;
        return this;
    }

    public Integer getThinkingBudget() {
        return thinkingBudget;
    }

    public AIChatEvent setThinkingBudget(Integer thinkingBudget) {
        this.thinkingBudget = thinkingBudget;
        return this;
    }

    public ThinkStatus getThinkStatus() {
        return thinkStatus;
    }

    public void setThinkStatus(ThinkStatus thinkStatus) {
        this.thinkStatus = thinkStatus;
    }

    public void appendResponse(String response){
        if (response == null) return;

        if(result==null){
            result = new AIMessage();
        }
        if (result.getContent() == null) {
            result.setContent(response);
        } else {
            result.setContent(result.getContent()+response);
        }
    }
    public void appendThink(String think){
        if (think == null) return;

        if(result==null){
            result = new AIMessage();
        }
        if (result.getThink() == null) {
            result.setThink(think);
        } else {
            result.setThink(result.getThink()+think);
        }
    }

    public void setResultId(String id) {
        if(result==null){
            result = new AIMessage();
        }
        result.setId(id);
    }

    public void setResultRole(AIMessageRole role) {
        if(result==null){
            result = new AIMessage();
        }
        result.setRole(role);
    }

    /**
     * 获取思考时长，单位ms
     * @return
     */
    public int getThinkDuration() {
        return (int)(thinkEndTime - thinkStartTime);
    }


    @Override
    protected void innerCancelHandle() {
        AIEventProcessor processor = this.getAttr(EventAttrKey.EVENT_PROCESSOR);
        AIResource resource = this.getAttr(EventAttrKey.EVENT_RESOURCE);
        if (processor != null) {
            processor.handleCancel(this, resource);
        }
    }
}
